import base64
from io import BytesIO
from docx import Document

# todo: 目前只是纯文本识别，若有 latex 或图片，还是需要调 ocr
# todo: 暂未处理数据字典，因为表头格式不统一，需要和产品协调
def get_docx_headings_and_formulas(dox_base64):
    try:
        doc_bytes = base64.b64decode(dox_base64)
        doc = Document(BytesIO(doc_bytes))
    except Exception as e:
        raise ValueError(f"Base64解码失败: {str(e)}")

    headings = []
    counters = [0] * 5
    current_formula = None
    collecting_formula = False

    for paragraph in doc.paragraphs:
        text = paragraph.text.strip()

        if collecting_formula:
            if "【说明】" in text or paragraph.style.name.startswith('Heading'):
                if headings:
                    headings[-1].update({
                        'has_formula': True,
                        'formula': current_formula,
                    })
                collecting_formula = False
                current_formula = None
            else:
                current_formula = f"{current_formula}\n{text}" if current_formula else text
                continue

        if paragraph.style.name.startswith('Heading') and text:
            level = int(paragraph.style.name.lstrip('Heading'))

            counters[level] += 1
            for l in range(level + 1, 5):
                counters[l] = 0

            headings.append({
                'level': '.'.join(str(counters[l]) for l in range(1, level + 1) if counters[l] > 0),
                'heading': paragraph.text,
                'has_formula': False,
                'formula': ''
            })

        elif '【公式】' in text:
            current_formula = ''
            collecting_formula = True

    if current_formula is not None and headings:
        headings[-1].update({
            'has_formula': True,
            'formula': current_formula,
        })

    return headings

import os
if __name__ == "__main__":
    home_dir = os.path.expanduser("~")
    # docx_path = os.path.join(home_dir, "Downloads", "文档导入-测试输入", "SOC文档.docx")
    docx_path = os.path.join(home_dir, "Downloads", "文档导入-测试输入", "DSP文档.docx")

    with open(docx_path, 'rb') as f:
        docx_base64 = base64.b64encode(f.read()).decode('utf-8')

    headings = get_docx_headings_and_formulas(docx_base64)
    for heading in headings:
        print(heading)